/*
 * License: Copyright (c) Huawei Technologies Co., Ltd. 2012-2019. All rights reserved.
 * Description: pull video stream module
 * Date: 2020-02-06 09:50:36
 * LastEditTime : 2020-02-19 20:37:04
 */

#include <sys/time.h>
#include <ctime>
#include "src_engine.h"
#include "hiaiengine/ai_memory.h"

const int FAILURE_NUM = 5;
const int SLEEP_TIME = 33;
const int SEND_TIMEOUT = 1000;
const int TIMES_MICROSECOND = 1000000;
const int MAX_QUEUE_NUM = 500;
const int SLEEP = 10;

uint64_t SrcEngine::GetCurentTimeStamp()
{
    struct timeval curTime = {
        .tv_sec = 0,
        .tv_usec = 0,
    };
    gettimeofday(&curTime, nullptr);
    return curTime.tv_sec * TIMES_MICROSECOND + curTime.tv_usec;
}

shared_ptr<AVFormatContext> SrcEngine::CreateFormatContext(string streamName)
{
    AVDictionary* options = nullptr;
    av_dict_set(&options, "rtsp_transport", "tcp", 0);
    av_dict_set(&options, "stimeout", "3000000", 0);
    AVFormatContext* formatContext = avformat_alloc_context();
    if (formatContext == NULL) {
        cout << "alloc context is failed!" << endl;
        return nullptr;
    }
    int ret = avformat_open_input(&formatContext, streamName.c_str(), nullptr, &options);
    if (options != nullptr) {
        av_dict_free(&options);
    }
    if (ret != 0) {
        avformat_free_context(formatContext);
        cout << "Couldn't open input stream " << streamName.c_str() << ",ret= " << ret << endl;
        return nullptr;
    }

    ret = avformat_find_stream_info(formatContext, nullptr);
    if (ret != 0) {
        avformat_free_context(formatContext);
        cout << "Couldn't find stream information." << endl;
        return nullptr;
    }
    return shared_ptr<AVFormatContext>(formatContext, [](AVFormatContext* p) {
        if (p) {
            avformat_close_input(&p);
        }
    });
}

int SrcEngine::GetStreamInfo(shared_ptr<AVFormatContext> pFormatCtx, int& format, int& videoIndex)
{
    if (pFormatCtx == nullptr) {
        cout << "pFormatCtx is null!" << endl;
        return -1;
    }

    videoIndex = -1;
    for (int i = 0; i < pFormatCtx->nb_streams; i++) {
        AVStream* inStream = pFormatCtx->streams[i];
        if (inStream->codecpar->codec_type != AVMEDIA_TYPE_VIDEO) {
            cout << "Error codec_type " << inStream->codecpar->codec_type << endl;
            return -1;
        }
        videoIndex = i;
        AVCodecID codecId = inStream->codecpar->codec_id;
        if (codecId == AV_CODEC_ID_H264) {
            format = 0;
        } else if (codecId == AV_CODEC_ID_H265) {
            format = 1;
        } else {
            cout << "Error unsupported format " << codecId << endl;
            return -1;
        }
    }

    if (videoIndex == -1) {
        cout << "Didn't find a video stream." << endl;
        return -1;
    }
    return 0;
}

HIAI_StatusT SrcEngine::SendPacketInfo(shared_ptr<TransData> output)
{
    int ret = HIAI_OK;
    do {
        ret = hiai::Engine::SendData(0, "TransData", static_pointer_cast<void>(output), SEND_TIMEOUT);
        if (ret == HIAI_QUEUE_FULL) {
            this_thread::sleep_for(chrono::milliseconds(SLEEP_TIME));
        }

        if (ret != HIAI_OK) {
            cout << "graphId: " << this->graphId << " deviceId: " << this->deviceId
                 << " StreamPuller send data failed. ErrCode: " << ret << endl;
            HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "StreamPuller send data failed. ErrCode: %x", ret);
        }
    } while (ret == HIAI_QUEUE_FULL);  // Wait for sending data, until trans_queue is not full.
    return ret;
}

void SrcEngine::TransferData()
{
    while (this->stop == 0) {
        mtx.lock();
        if (dataQueues.empty()) {
            mtx.unlock();
            this_thread::sleep_for(chrono::milliseconds(SLEEP));
            continue;
        }
        shared_ptr<PacketInfo> sendPacket = dataQueues.front();
        dataQueues.pop();
        mtx.unlock();

        shared_ptr<TransData> output = make_shared<TransData>();
        output->graphId = this->graphId;
        output->deviceId = this->deviceId;
        output->channelId = this->channelId;
        output->outputFormat = this->mediaAttr.outputFormat;
        output->vencFrameRate = this->mediaAttr.vencFrameRate;
        output->vencBitRate = this->mediaAttr.vencBitRate;
        output->videoFormat = this->mediaAttr.videoFormat;
        output->vencCtlType = this->mediaAttr.vencCtlType;
        output->eos = 0;
        output->buf = sendPacket->buf;
        output->len = sendPacket->len;
        output->frameId = sendPacket->frameNum;
        output->streamPullFps = sendPacket->fps;

        SendPacketInfo(output);
    }

    int sendFailNum = 0;
    shared_ptr<TransData> stopInfo = make_shared<TransData>();
    stopInfo->graphId = this->graphId;
    stopInfo->deviceId = this->deviceId;
    stopInfo->channelId = this->channelId;
    stopInfo->eos = 1;
    stopInfo->len = 0;
    do {
        if (SendPacketInfo(stopInfo) == HIAI_OK) {
            break;
        }
        sendFailNum++;
    } while (sendFailNum < FAILURE_NUM);
}

int SrcEngine::ReadFrame(shared_ptr<AVFormatContext> pFormatCtx, uint64_t& frameNum, int videoIndex)
{
    AVPacket pkt;
    while (true) {
        av_init_packet(&pkt);
        int ret = av_read_frame(pFormatCtx.get(), &pkt);
        if (ret == AVERROR_EOF) {
            cout << " End of file. GraphId:" << this->graphId << " deviceId: " << this->deviceId
                 << "Total frame number: " << frameNum << endl;
            break;
        } else if (ret != 0) {
            cout << "graphId: " << this->graphId << " deviceId: " << this->deviceId
                 << " Read frame failed. ErrCode: " << ret << endl;
            continue;
        } else if (pkt.stream_index == videoIndex) {
            if (pkt.size <= 0) {
                cout << "graphId " << this->graphId << " deviceId: " << this->deviceId
                     << " Invalid pkt.size= " << pkt.size << endl;
                continue;
            }

            uint8_t* buffer = new uint8_t[pkt.size];
            if (buffer == nullptr) {
                cout << "graphId: " << this->graphId << " deviceId: " << this->deviceId << " new memory faild." << endl;
                av_packet_unref(&pkt);
                continue;
            }
            ret = memcpy_s(buffer, pkt.size, pkt.data, pkt.size);
            if (ret != 0) {
                cout << "graphId: " << this->graphId << " deviceId: " << this->deviceId << " memcpy faild." << endl;
                delete[] buffer;
                av_packet_unref(&pkt);
                continue;
            }
            shared_ptr<PacketInfo> packetInfo = make_shared<PacketInfo>();
            packetInfo->buf.reset(buffer, [](uint8_t* p) { delete[] p; });
            packetInfo->len = pkt.size;
            frameNum += 1;
            packetInfo->frameNum = frameNum;
            packetInfo->fps = frameNum / (1.0 * (this->GetCurentTimeStamp() - this->initTimeStamp) / TIMES_MICROSECOND);
            mtx.lock();
            if (dataQueues.size() < MAX_QUEUE_NUM) {
                dataQueues.push(packetInfo);
            }
            mtx.unlock();
        }
        av_packet_unref(&pkt);
    }
    return -1;
}

void SrcEngine::PullStreamDataLoop()
{
    int failNum = 0;
    uint64_t frameNum = 0;
    shared_ptr<AVFormatContext> pFormatCtx;
    while (failNum < FAILURE_NUM) {
        pFormatCtx = CreateFormatContext(g_videos[this->channelId]);
        if (pFormatCtx == nullptr) {
            cout << "Create context failed. graphId: " << this->graphId << " deviceId: " << this->deviceId
                 << " fail times: " << ++failNum << endl;
            this_thread::sleep_for(chrono::milliseconds(SLEEP_TIME));
            continue;
        }

        // get stream infomation
        int format = 0;
        int videoIndex = -1;
        if (GetStreamInfo(pFormatCtx, format, videoIndex) == -1) {
            cout << "Get stream failed. graphId: " << this->graphId << " deviceId: " << this->deviceId
                 << " fail times: " << ++failNum << endl;
            continue;
        }
        this->mediaAttr.videoFormat = format;

        int ret = ReadFrame(pFormatCtx, frameNum, videoIndex);
        cout << "Pull failed. graphId: " << this->graphId << " deviceId: " << this->deviceId
             << " fail times: " << ++failNum << endl;
    }
    this->stop = 1;
    cout << "Pull stream failure. graphId: " << this->graphId << " deviceId: " << this->deviceId << endl;
}

HIAI_StatusT SrcEngine::StartStream()
{
    pullStreamThread = thread(&SrcEngine::PullStreamDataLoop, this);
    sendDataThread = thread(&SrcEngine::TransferData, this);
    return HIAI_OK;
}

HIAI_StatusT SrcEngine::Init(const hiai::AIConfig& config, const vector<hiai::AIModelDescription>& model_desc)
{
    avformat_network_init();
    return HIAI_OK;
}

SrcEngine::~SrcEngine()
{
    avformat_network_deinit();
    if (pullStreamThread.joinable()) {
        pullStreamThread.join();
    }

    if (sendDataThread.joinable()) {
        sendDataThread.join();
    }
}

HIAI_IMPL_ENGINE_PROCESS("SrcEngine", SrcEngine, RP_INPUT_SIZE)
{
    if (arg0 != nullptr) {
        if (initTimeStamp == 0) {
            initTimeStamp = this->GetCurentTimeStamp();
        }
        shared_ptr<TransData> input = static_pointer_cast<TransData>(arg0);
        this->graphId = input->graphId;
        this->deviceId = input->deviceId;
        this->channelId = input->channelId;
        this->mediaAttr.SetAttr(-1, input->outputFormat, input->vencFrameRate, input->vencBitRate, input->vencCtlType);
        StartStream();
    }
    return HIAI_OK;
}